package com.zmc.dynamic.shard.interceptor;

import com.baomidou.mybatisplus.annotation.TableName;
import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import com.baomidou.mybatisplus.core.toolkit.StringPool;
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.reflection.DefaultReflectorFactory;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.apache.shardingsphere.sharding.yaml.config.rule.YamlTableRuleConfiguration;
import org.apache.shardingsphere.sharding.yaml.config.strategy.sharding.YamlShardingStrategyConfiguration;

import java.sql.Connection;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Slf4j
public class MyBatisPlusShardingInterceptor implements InnerInterceptor {

    private Map<String, YamlTableRuleConfiguration> shardTableMap;

    public MyBatisPlusShardingInterceptor(Map<String, YamlTableRuleConfiguration> shardTableMap) {
        this.shardTableMap = shardTableMap;
    }

    @Override
    public void beforePrepare(StatementHandler sh, Connection connection, Integer transactionTimeout) {
        prepare(sh);
    }

    /**
     * 拦截Sql语法构建的处理
     * 处理sql的分表属性拼接
     */
    private void prepare(StatementHandler statementHandler) {
        log.debug("prepare");
        // 通过MetaObject优雅访问对象的属性，这里是访问statementHandler的属性;：MetaObject是Mybatis提供的一个用于方便、
        // 优雅访问对象属性的对象，通过它可以简化代码、不需要try/catch各种reflect异常，
        // 同时它支持对JavaBean、Collection、Map三种类型对象的操作。
        MetaObject metaObject = MetaObject.forObject(statementHandler, SystemMetaObject.DEFAULT_OBJECT_FACTORY,
                SystemMetaObject.DEFAULT_OBJECT_WRAPPER_FACTORY, new DefaultReflectorFactory());
        // 获取mapper的Statement对象,它描述的是mapper对象的配置
        MappedStatement mappedStatement = (MappedStatement) metaObject.getValue("delegate.mappedStatement");
        // 判断是否为更新方法
        if (SqlCommandType.UPDATE != mappedStatement.getSqlCommandType()) {
            return;
        }
        String msId = mappedStatement.getId();
        String methodName = msId.substring(msId.lastIndexOf(StringPool.DOT) + 1);
        if (!"updateById".equals(methodName)) {
            log.debug("prepare methodName:{} out", methodName);
            return;
        }
        // 获取数据库对象,此处需要@TableName注解获取表名并进行分表表名匹配
        Class<?> pojoClazz = mappedStatement.getParameterMap().getType();
        if (!pojoClazz.isAnnotationPresent(TableName.class)) {
            log.debug("prepare {} 缺少 @TableName", pojoClazz.getName());
            return;
        }
        TableName annotation = pojoClazz.getAnnotation(TableName.class);
        if (!shardTableMap.containsKey(annotation.value())) {
            log.debug("prepare {} 缺少 @TableName", pojoClazz.getName());
            return;
        }
        // 在sql尾部拼接分表所属属性
        YamlTableRuleConfiguration shardConfig = shardTableMap.get(annotation.value());
        YamlShardingStrategyConfiguration tableStrategy = shardConfig.getTableStrategy();
        List<String> list = new ArrayList<>();
        if (tableStrategy.getStandard() != null) {
            String shardingColumn = tableStrategy.getStandard().getShardingColumn();
            if (shardingColumn != null) {
                list.add(shardingColumn);
            }
        } else if (tableStrategy.getComplex() != null) {
            List<String> collect = Arrays.stream(tableStrategy.getComplex().getShardingColumns().split(",")).collect(Collectors.toList());
            if (CollectionUtils.isNotEmpty(collect)) {
                list.addAll(collect);
            }
        }
        BoundSql boundSql = (BoundSql) metaObject.getValue("delegate.boundSql");
        String sql = boundSql.getSql();
        String[] newSql = {sql};
        list.stream().filter(column -> !column.equals("id")).forEach(shardingColumn -> {
            newSql[0] = newSql[0].replace((shardingColumn + "=?,"), "");
        });
        metaObject.setValue("delegate.boundSql.sql", newSql[0]);
    }

}
